#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：div-align-dg 
@File    ：visualize_augs.py
@IDE     ：PyCharm 
@Author  ：cao xu
@Date    ：2025/9/5 上午11:10 
"""
import os
import sys
import math
import random
import argparse
from typing import List, Tuple, Callable, Optional
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from imagecorruptions import corrupt

AUG_NAMES = [
    "defocus_blur",
    "glass_blur",
    "gaussian_blur",
    "motion_blur",
    "speckle_noise",
    "shot_noise",
    "impulse_noise",
    "gaussian_noise",
    "jpeg_compression",
    "pixelate",
    "elastic_transform",
    "brightness",
    "saturate",
    "contrast",
    "high_pass_filter",
    # "phase_scaling",
]

CN_LABELS = {
    "original": "原图",
    "defocus_blur": "失焦模糊",
    "glass_blur": "玻璃模糊",
    "gaussian_blur": "高斯模糊",
    "motion_blur": "运动模糊",
    "speckle_noise": "斑点噪声",
    "shot_noise": "散粒噪声",
    "impulse_noise": "脉冲噪声",
    "gaussian_noise": "高斯噪声",
    "jpeg_compression": "JPEG压缩",
    "pixelate": "像素化",
    "elastic_transform": "弹性形变",
    "brightness": "亮度",
    "saturate": "饱和度",
    "contrast": "对比度",
    "high_pass_filter": "高通滤波",
    # "phase_scaling": "相位缩放",
}

# ============== 尝试复用你项目里的训练时增强（优先） ==============
def build_aug_from_project(project_root: Optional[str]) -> Optional[Callable[[np.ndarray, str, int], np.ndarray]]:
    if not project_root:
        return None
    if project_root not in sys.path:
        sys.path.insert(0, project_root)
    # 下面两种导入方式，按你的仓库实际路径择一；都失败则返回 None
    try:
        from maskrcnn_benchmark.data.transforms.corruptions import corrupt as project_corrupt  # type: ignore
        def _fn(img_np: np.ndarray, name: str, severity: int) -> np.ndarray:
            return project_corrupt(img_np, name, severity)
        return _fn
    except Exception:
        pass
    try:
        from maskrcnn_benchmark.data.transforms.corruptions import get_corruption_fn  # type: ignore
        def _fn(img_np: np.ndarray, name: str, severity: int) -> np.ndarray:
            fn = get_corruption_fn(name)
            return fn(img_np, severity)
        return _fn
    except Exception:
        pass
    return None

def build_aug_from_imagecorruptions() -> Callable[[np.ndarray, str, int], np.ndarray]:
    def _fn(img_np: np.ndarray, name: str, severity: int) -> np.ndarray:
        return corrupt(img_np, corruption_name=name, severity=severity)
    return _fn

# ============== 两个少见增强的占位：若你项目有真实实现，请在上面导入并替换掉这里 ==============
def fallback_high_pass_filter(img_np: np.ndarray, severity: int) -> np.ndarray:
    try:
        from scipy.ndimage import gaussian_filter
    except Exception:
        raise RuntimeError("需要 scipy：pip install scipy")
    img = img_np.astype(np.float32)
    sigma_map = {1: 1.0, 2: 1.5, 3: 2.0, 4: 3.0, 5: 4.0}
    sigma = sigma_map.get(int(severity), 2.0)
    low = gaussian_filter(img, sigma=(sigma, sigma, 0))
    high = img - low + 128.0
    return np.clip(high, 0, 255).astype(np.uint8)

def fallback_phase_scaling(img_np: np.ndarray, severity: int) -> np.ndarray:
    scale_map = {1: 0.8, 2: 0.6, 3: 0.4, 4: 0.3, 5: 0.2}
    scale = scale_map.get(int(severity), 0.4)
    out = []
    for c in range(3):
        x = img_np[..., c].astype(np.float32)
        F = np.fft.fft2(x)
        mag, ph = np.abs(F), np.angle(F)
        F2 = mag * np.exp(1j * (ph * scale))
        x2 = np.fft.ifft2(F2).real
        x2 = x2 - x2.min()
        x2 = x2 / (x2.max() + 1e-8) * 255.0
        out.append(x2.astype(np.uint8))
    return np.stack(out, axis=-1)

# ============== 绘制网格 & 标签 ==============
def draw_label(pil_img: Image.Image, text: str, pad: int = 6) -> Image.Image:
    draw = ImageDraw.Draw(pil_img)
    try:
        font = ImageFont.truetype("NotoSansCJK-Regular.ttc", 22)
    except Exception:
        font = ImageFont.load_default()
    tw, th = draw.textsize(text, font=font)
    rect = [0, 0, tw + pad * 2, th + pad * 2]
    draw.rectangle(rect, fill=(0, 0, 0, 160))
    draw.text((pad, pad), text, fill=(255, 255, 255), font=font)
    return pil_img

def make_grid(tiles: List[Tuple[Image.Image, str]], cell_wh=(384, 256)) -> Image.Image:
    W, H = cell_wh
    n = len(tiles)
    if n == 16:  # 原图 + 15 种增强
        g = 4
    else:
        g = math.ceil(math.sqrt(n))
    canvas = Image.new("RGB", (g * W, g * H), color=(255, 255, 255))
    for i, (im, label) in enumerate(tiles):
        r, c = divmod(i, g)
        img = im.resize((W, H), Image.BILINEAR)
        img = draw_label(img, label)
        canvas.paste(img, (c * W, r * H))
    return canvas

# ============== 文件收集 ==============
IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}

def list_images_recursive(root: str) -> List[str]:
    files = []
    for d, _, fs in os.walk(root):
        for f in fs:
            if os.path.splitext(f.lower())[1] in IMG_EXTS:
                files.append(os.path.join(d, f))
    return files

# ============== 主流程 ==============
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-root", required=True, help="医院目录的上一级：每个子目录是一个医院")
    ap.add_argument("--out-root", default="./aug_vis_out", help="输出根路径")
    ap.add_argument("--per-hospital", type=int, default=10, help="每家医院采样的图片数")
    ap.add_argument("--severity", type=int, default=3, help="增强严重度（与训练一致）")
    ap.add_argument("--project-root", type=str, default="", help="你的仓库根目录（若有自定义增强）")
    ap.add_argument("--seed", type=int, default=123, help="随机种子")
    ap.add_argument("--cell-width", type=int, default=384)
    ap.add_argument("--cell-height", type=int, default=256)
    args = ap.parse_args()

    random.seed(args.seed)
    os.makedirs(args.out_root, exist_ok=True)

    # 选择增强实现
    aug_fn_project = build_aug_from_project(args.project_root)
    if aug_fn_project is None:
        try:
            aug_fn = build_aug_from_imagecorruptions()
            using_project = False
            print("[Info] 使用 imagecorruptions 实现增强（与常见训练配置一致）。")
        except Exception as e:
            raise RuntimeError("找不到增强实现：请提供 --project-root 或安装 imagecorruptions") from e
    else:
        aug_fn = aug_fn_project
        using_project = True
        print("[Info] 使用项目中的增强实现（与训练完全一致）。")

    hospitals = sorted([d for d in os.listdir(args.data_root) if os.path.isdir(os.path.join(args.data_root, d))])
    print(f"[Info] 发现医院目录 {len(hospitals)} 个。")

    for hid, hosp in enumerate(hospitals, 1):
        hosp_dir = os.path.join(args.data_root, hosp)
        out_dir = os.path.join(args.out_root, hosp)
        os.makedirs(out_dir, exist_ok=True)

        imgs = list_images_recursive(hosp_dir)
        if not imgs:
            print(f"[Warn] 医院 {hosp} 下未找到图片，跳过。")
            continue
        sample = random.sample(imgs, k=min(args.per_hospital, len(imgs)))

        print(f"[{hid}/{len(hospitals)}] {hosp}: 采样 {len(sample)} 张")

        for idx, img_path in enumerate(sample):
            try:
                img = Image.open(img_path).convert("RGB")
                npimg = np.array(img)

                tiles = [(img.copy(), CN_LABELS["original"])]
                for name in AUG_NAMES:
                    if name == "high_pass_filter":
                        out_np = fallback_high_pass_filter(npimg, args.severity)
                    elif name == "phase_scaling":
                        out_np = fallback_phase_scaling(npimg, args.severity)
                    else:
                        out_np = aug_fn(npimg, name, args.severity)
                    tiles.append((Image.fromarray(out_np), CN_LABELS.get(name, name)))

                grid = make_grid(tiles, (args.cell_width, args.cell_height))
                # 输出文件名：避免中文问题，统一用英文+索引
                base = f"hospital_{hid:02d}_{hosp}_idx{idx:02d}_grid.png"
                # 去掉文件名中的空格或可能不兼容字符
                safe_base = "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in base)
                save_path = os.path.join(out_dir, safe_base)
                grid.save(save_path)
            except Exception as e:
                print(f"[Warn] 处理 {img_path} 失败：{e}")

    print("[Done] 全部完成。输出在：", args.out_root)


if __name__ == "__main__":
    main()
